import numpy as np
import cv2 as cv
import matplotlib.pyplot as plt
import math
import axis_transform
'''
cv库中的矩阵运算采用32位->8位有效数字
因此，在卡尔曼融合运算中只进行了小数运算
'''
mean = np.array([2, 1])
conv = np.array([[0.5, 0.0], [0.0, 0.5]])
x, y = axis_transform.get_xy(0)
x, y = x - 118, y - 31
last_measurement = current_measurement = np.array([[np.float32(x)], [np.float32(y)]], np.float32)
last_prediction = current_prediction = np.array([[np.float32(x)], [np.float32(y)]], np.float32)

kalman = cv.KalmanFilter(2, 2)
kalman.measurementMatrix = np.array([[1, 0], [0, 1]], np.float32)
kalman.transitionMatrix = np.array([[1, 0], [0, 1]], np.float32)
kalman.processNoiseCov = np.array([[0.003, 0.0], [0.0, 0.003]], np.float32)


# kalman.statePost = np.array([[np.float32(x)], [np.float32(y)]], np.float32)


def normal_number(x, y):
    global current_measurement, last_measurement, current_prediction, last_prediction
    last_prediction = current_prediction  # 把当前预测存储为上一次预测
    last_measurement = current_measurement  # 把当前测量存储为上一次测量
    current_measurement = np.array([[np.float32(x)], [np.float32(y)]])  # 当前测量
    current_correct = kalman.correct(current_measurement)  # 用当前测量来校正卡尔曼滤波器
    print("current_correct:", current_correct)
    print("gain:", kalman.gain)
    current_prediction = kalman.predict()  # 计算卡尔曼预测值，作为当前预测
    print("current_prediction:", current_prediction)


def monteCarlo():
    for i in range(1000):
        x, y = np.random.multivariate_normal(mean=mean, cov=conv, size=1).T
        normal_number(x, y)
        cpx, cpy = current_prediction[0], current_prediction[1]
        print(cpx, cpy)
        plt.plot(x, y, 'xr')  # x为×，r为红色

    plt.plot(cpx, cpy, 'go')
    plt.show()


'''
卡尔曼滤波核心代码：
输入起止时间，每秒一帧计算
'''


def run_kalman_position(start_time, end_time):
    start = math.floor(start_time)
    end = math.ceil(end_time)
    prediction_x, prediction_y = 0.0, 0.0
    plt.figure()
    plt.xlabel("Longitude")
    plt.ylabel("Latitude")
    plt.grid()
    for i in range(end - start):
        cur_time = i + start
        x, y = axis_transform.get_xy(cur_time)
        x, y = x - 118, y - 31
        if cur_time == 0:
            kalman.statePre = np.array([[x], [y]], np.float32)
            kalman.errorCovPre = np.array([[0.8, 0.0], [0.0, 0.8]], np.float32)
            continue
        if cur_time in [1, 2, 21, 27, 28]:  # 移除误差较大的点
            continue
        normal_number(x, y)
        prediction_x, prediction_y = current_prediction[0], current_prediction[1]
        point1, = plt.plot(x + 118, y + 31, 'rx')  # x为×，r为红色

    point2, = plt.plot(prediction_x + 118, prediction_y + 31, 'go')
    point3, = plt.plot([118.790087], [31.939067], 'bo')  # 真实位置
    plt.legend(handles=[point2, point1, point3], labels=['prediction position', 'measurement position',
                                                         'target position'])

    plt.show()


if __name__ == "__main__":
    run_kalman_position(0, 29)
